#include "hw4.hh"

/***********************************************************
 * parComputeMorton
 *   Compute morton IDs by discretizing domain
 *
 *   x - set of points in domain
 *   N - number of points
 *   b - number of buckets
 ***********************************************************/
double parComputeMorton ( Point *x, int N, int b ) {
	double min = x[0].x;					  // min x coord
	double h = (x[N-1].x - min) / (double) b; // bucket size
	int i;

	  // h is bucket size. 
	  // Add tiny amount to get morton IDs in [0, b-1)
	h += h * 0.0001;

	#pragma omp parallel for
	for (i = 0; i < N; i++)	x[i].morton = floor((x[i].x - min) / h); 

	return h;
} // parComputemorton

inline long int morton ( double x, double min, double h ) {
	return floor( x - min ) / h;
}

/***********************************************************
 * parComputeDepth
 *   Compute depth by finding min gap between seq elements
 *   and dividing space by that gap
 *
 *   x - set of points
 *   N - number of points
 *
 *   return - depth of tree
 ***********************************************************/
int parComputeDepth ( double *x, int N ) {
	int i;			// iterator
	int b;			// buckets
	double h = 99;	// smallest gap size so far
	double d;		// current gap size

	#pragma omp parallel for private(d) 
	for (i = 0; i < N-1; i++) {
		d = x[i+1] - x[i];
		  // shared across threads. be careful.
		#pragma omp critical 
		if ( d < h ) { 
			h = d;
		}
	}

	  // Slightly reduce gap to get correct number of buckets
	h *= 0.9999;
	b = pow2roundup((int)((x[N-1] - x[0]) / h));

	  // return log2 of number of buckets
	return (int)ceil(log2(b));
} // parComputeDepth


/***********************************************************
 * G
 *   Exact evaluation kernel
 *
 *   y - array of source points
 *   d - densities of source points
 *   N - number of source points
 *   x - array of target points
 *   u - array of sums of forces on target points
 *   M - number of target points
 ***********************************************************/
void G ( double *y, double *d, int N, double *x, double *u, int M ) {
	int i, j;				// iterators
	double t;				// temporary aggregation of forces

	for (i = 0; i < M; i++) {
		t = 0;
		for (j = 0; j < N; j++) {
			if (x[i] != y[j]) 
				t += d[j]/(x[i]-y[j]);
		}
		u[i] = t;
	}
} // G


/***********************************************************
 * TreeMerge 
 *   Merge T2 into T1
 *
 *   T1 - first tree to merge, also destination
 *   T2 - second tree to merge
 ***********************************************************/
void TreeMerge ( TreePtr *T1, TreePtr *T2 ) {
	MapIter it, it2;	// iterators
	Point p, p2;		// temp points

	  // merge T2 with T1 in key order
	T1->M.insert(T2->M.begin(), T2->M.end());

	  // if only one element, nothing to do. move along.
	if (T1->M.size() < 2) return;

	  // initialize our iterators
	  //   it = first element. it2 = second element
	it = T1->M.begin(); it2 = it; it2++;	

	  // loop until it2 exits the map
	while (it2 != T1->M.end()) {
		  // if we have a duplicate
		if ((*it).first == (*it2).first) {
			  // store point info in temp vars
			p = *((*it).second);
			p2 = *((*it2).second);

			if (!(p.isLeaf || p2.isLeaf)) {
			  	  // Both internal, just delete one element
				T1->M.erase(it2++);
			} else if (p.isLeaf && p2.isLeaf) {
				  // Both leaf, delete and insert one element
				  //   Insert will push other element down tree
				  //   Decrement iterator out of the danger zone
				  //     to keep it valid
				it--;
				T1->M.erase(it2);
				T1->Insert(p2);
				it++; it2 = it; it2++;
			} else {
				  // Mixed leaf/internal, delete the leaf and reinsert
				if (p.isLeaf) {
					T1->M.erase(it++);
					T1->Insert(p);
					it2++;
				} else {
					T1->M.erase(it2);
					T1->Insert(p2);
					it2 = it; it2++;
				}
			}	
		} else {
			 // No duplicate, move along
			it++; it2++;
		}
	}
} // TreeMerge


/***********************************************************
 * wellSeparated 
 *    Determines whether morton IDs are well-separated
 *    based on level and depth of tree. Well-separated is
 *    defined as the target node (m2) not being a descendent
 *    of the source's (m1's) neighbors or self.
 *
 *    m1 = nodeID of node in tree
 *    m2 = morton ID of target point
 *    lev = level of node in tree
 *    d = depth of tree
 *
 *    return - 1 if well-separated, 0 if not
 ***********************************************************/
int wellSeparated( long int m1, long int m2, int lev, int d ) {
	int shift = d - lev;
	long int diff;

	if (shift < 0) cout << "BAD SHIFT" << endl;
	m1 >>= shift;
	m2 >>= shift;
	diff = m1 - m2;

	return (diff < -1 || diff > 1) ? 1 : 0;
} // wellSeparated


/***********************************************************
 * evaluateHelper
 *    Single-point, recursive evaluation
 *
 *    T - tree containing y 
 *    y - source point
 *    x - target point 
 *
 *    return - sum of all forces exerted on m by p and its children
 ***********************************************************/
double evaluateHelper ( TreePtr *T, Point y, Point x) {
	double u, uL, uR;

	u = uL = uR = 0.0;

	if ( y.isLeaf || wellSeparated( y.nodeID, x.morton, y.level, T->depth ) ) {
		T->G( &(y.x), &(y.density), 1, &(x.x), &u, 1 );
		  // evaluate at p 
		return u;
	} else {
		#pragma omp parallel sections
		{
			  // evaluate left
			#pragma omp section
			if ( T->M.count(y.Left()) != 0 ) 
				uL = evaluateHelper( T, T->PointFromKey(y.Left()), x );
			
			  // evaluate right
			#pragma omp section
			if ( T->M.count(y.Right()) != 0 ) 
				uR = evaluateHelper( T, T->PointFromKey(y.Right()), x );
		}

		return uL + uR;
	}
} // evaluateHelper


/***********************************************************
 * evaluate
 *    Evaluate all target points against source tree
 *
 *    T - tree against which to evaluate points
 *    x - target points
 *    u - calculated forces
 *    M - # of points
 ***********************************************************/
void evaluate ( TreePtr *T, double *x, double *u, int M ) {
	int i;				// iterator
	long int m;			// morton ID of current point
	Point x_pt;			// Point representation of target point

	x_pt = Point();

	for (i = 0; i < M; i++) {
		m = morton(x[i], T->y[0].x, T->h);
		x_pt.x = x[i]; 
		x_pt.morton = m;
		u[i] = evaluateHelper (T, T->PointFromKey( 0 ), x_pt);
	}
	
} // evaluate


/***********************************************************
 * setup
 *   Core of the program, drives the algorithm.
 *
 *   y - set of points in domain
 *   d - densities of points
 *   N - number of points
 *   G - pointer to computation kernel
 *
 *   return - pointer to complete tree
 ***********************************************************/
TreePtr *setup ( double *y, double *d, int N, 
		void (*G)(double *y, double *d, int N, double *x, double *u, int M)) {
	int i, j, k;				// Iterators.
	int logThreads;				// Log2 of numThreads.
	int depth;					// Depth of tree.
	int ppt;					// Points per thread.
	int thr;					// Number of running thread
	double h;					// size of morton buckets
	TreePtr *T;					// Main btree of Points.
	TreePtr *TA[numThreads];	// Array of one tree per thread.
	Point *y_pt;				// Array of Points.

	  // sort points in coordinate order
	parqsort<double>(y, 0, N-1, &cmpDouble, &swapDouble);

	  // calculate depth of tree
	depth = parComputeDepth( y, N );

	  // calculate points per thread
	ppt = ceil( 1.0 * N / numThreads );

	  // allocate space for array of points
	y_pt = new Point[N];

	  // initialize array of points
	for (i = 0; i < N; i++) y_pt[i] = Point( y[i], d[i] );
	
	  // sort points and compute Morton IDs in parallel
	h = parComputeMorton(y_pt, N, 1 << depth);

	  // Initialize trees
	for (i = 0; i < numThreads; i++) TA[i] = new TreePtr ( depth );
	
	  // insert points into trees
	#pragma omp parallel for schedule(static, ppt) private(thr) 
	for (i = 0; i < N; i++) {
		thr = omp_get_thread_num();
		TA[thr]->Insert( y_pt[i] );
	}

	cout << "Insert done" << endl;

 	  // Parallel tree merge
	  //   pair off trees by dimension
	logThreads = (int)ceil(log2(numThreads));
	for (i = logThreads-1; i >= 0; i--) {
		k = 1 << i;
		#pragma omp parallel for 
		for (j = 0; j < k; j++) 
			if (j + k < numThreads) 
				TreeMerge(TA[j], TA[j + k]);
	}
	
	cout << "Merge done" << endl;

	T = TA[0]; 						// T = full tree
	T->G = G;						// set computation kernel
	T->y = y_pt;					// keep original set of points
	T->N = N;						// set number of points in tree
	T->h = h;						// set bucket size in tree
	T->ppt = ppt;					// set points per thread
	T->BuildVector();				// compute vector representation
	T->BuildLeafVector();			// compute leaf vector representation
	T->parComputeSubtreeSizes();	// parallel compute subtree sizes
	T->parComputeEulerTour();		// parallel compute Euler tour
	T->ComputeAverage();			// compute averages/sums in tree
	
	return T;
} // setup


/***********************************************************
 * main 
 *   Read params and set up points, set omp options.
 ***********************************************************/
int main ( int argc, char **argv ) {
	int N = 20;					// # of points. Default to 20.
	int i;						// iterator
	double *y;					// array of points
	double *d;					// array of densities
	double *u;					// sum of forces on each point
	double t1, t2;
	TreePtr *T;

	  // Read in N and numThreads if given.
	if (argc > 1) N = atoi(argv[1]);
	if (argc > 2) numThreads = atoi(argv[2]);

	  // Allocate space for array of Points.
	y = new double[N];
	d = new double[N];

	  // Seed random number generator
	srand48( time(NULL) );

	  // Seed first half of array w/normal dist: u=0.3, v=0.05
	for (i = 0; i < N/2; i++) y[i] = normal(0.3, 0.05);	

	  // Seed second half of array w/normal dist: u=0.6, v=0.01
	for (i = N/2; i < N; i++) y[i] = normal(0.6, 0.01);

	  // Seed densities
	for (i = 0; i < N/2; i++) d[i] = normal(100, 20);
	for (i = N/2; i < N; i++) d[i] = normal(30, 10);

	  // Set number of threads, allow dynamic and nested
	omp_set_dynamic( 1 );
//	omp_set_nested( 1 );
	omp_set_num_threads( numThreads );

	u = new double[N];

	t1 = omp_get_wtime();

	  // Set-up tree
	T = setup( y, d, N, &G );
	cout << "Done setup" << endl;

	  // Evaluate target points against tree T
	evaluate( T, y, u, N );

	t2 = omp_get_wtime();

	cout << "Took " << t2-t1 << " seconds." << endl;
	return 0;
} // main


